/**
 * Note:
 * p: parent
 * g: grandparent
 * u: uncle
 */
package com.beney.ds.trees.rbTree;

import java.util.LinkedList;
import java.util.Queue;

public class RBTree<K extends Comparable<K>, V> {
  private static final boolean RED = true;
  private static final boolean BLACK = false;

  Node root;  // rbtree 的根

  private class Node {
    private K key;            // 可比较的关键码
    private V val;          // 结点的数据
    private Node parent, left, right;   // 父\左\右子节点指针
    private boolean color;      // 结点颜色
    private int size;           // 当前结点的子树大小

    public Node(K key, V val, boolean color, int size) {
      this.key = key;
      this.val = val;
      this.color = color;
      this.size = size;
    }
  }

  public RBTree() {
  }

  /******************************************************
   *  Node helper methods
   ******************************************************/

  private int size(Node n) {
    if (n == null) return 0;
    return n.size;
  }

  public int size() {
    return size(root);
  }

  public boolean isEmpty() {
    return root == null;
  }

  /*****************************************************
   *  Standard BST search
   *****************************************************/

  public V get(K key) {
    if (key == null) throw new IllegalArgumentException("get(key), key==null");
    return get(root, key);
  }

  public boolean contains(K key) {
    return get(key) != null;
  }

  private V get(Node n, K key) {
    while (n != null) {
      int cmp = key.compareTo(n.key);
      if (cmp < 0) n = n.left; // 左小右大
      else if (cmp > 0) n = n.right;
      else return n.val;
    }
    return null;
  }

  /******************************************************
   *  RBTree insertion: bst insertion + fix insertion
   ******************************************************/

  public void put(K key, V val) {
    if (key == null) throw new IllegalArgumentException("put(key, val), key==null");
    if (val == null) {
      delete(key);
    }
    put(root, key, val);
  }


  private void put(Node n, K key, V val) {
    if (root == null) {
      root = new Node(key, val, BLACK, 1);
      return;
    }
    Node p = n.parent;
    int util = 0;
    while (n != null) {
      if ((util = key.compareTo(n.key)) < 0) {
        p = n;
        n = n.left;
      } else if (util > 0) {
        p = n;
        n = n.right;
      } else {
        n.val = val;
        return;
      }
    }
    n = new Node(key, val, RED, 1);
    n.parent = p;
    if (util > 0) p.right = n;
    else p.left = n;
    fixPut(n);
  }

  private void fixPut(Node k) {
    Node u, g, p;   // uncle, grandparent, parent
    // case2: parent is BLACK, didn't violate any property
    while ((p = k.parent) != null && p.color == RED) {  // case3: parent is red, must have grandparent
      g = p.parent;
      if (p == g.right) {
        u = g.left;
        if (u != null && u.color == RED) {          // case3.1
          flipColors(g);
          k = g;                                  // only in case3.1, we have g shift to RED that should be cared
        } else {                // uncle.color = BLACK
          if (k == p.left) {  // case3.2.2: right left case
            k = p;
            p = p.left;
            rotateRight(k);
          }                   // case3.2.1: right right case
          g.color = !g.color;
          p.color = !p.color;
          rotateLeft(g);
        }
      } else if (p == g.left) {
        u = g.right;
        if (u != null && u.color == RED) {   // case3.1
          flipColors(g);
          k = g;
        } else {                // uncle.color = BLACK
          if (k == p.right) { // case3.2.4: left right case
            k = p;
            p = p.right;
            rotateLeft(k);
          }                   // case3.2.3: left left case
          g.color = !g.color;
          p.color = !p.color;
          rotateRight(g);
        }
      }
    }
    root.color = BLACK;
  }


  /****************************************************
   *  RBTree deletion
   ****************************************************/

  private void delete(K key) {
  }

  /***************************************************
   *  RBTree helper functions
   ***************************************************/

  /**
   * do right rotation at n
   *
   * @param n the pivot of rotation
   */
  private void rotateRight(Node n) {
    Node l = n.left;
    n.left = l.right;
    l.right.parent = n;
    l.right = n;

    if (n == root) root = l;
    else if (n.parent.left == n) {
      n.parent.left = l;
    } else {
      n.parent.right = l;
    }

    l.parent = n.parent;
    n.parent = l;

    l.size = n.size;
    n.size = size(n.left) + size(n.right) + 1;
  }

  /**
   * do left rotation at n
   *
   * @param n the pivot of rotation
   */
  private void rotateLeft(Node n) {
    Node r = n.right;
    n.right = r.left;
    r.left.parent = n;
    r.left = n;

    if (n == root) root = r;
    else if (n.parent.left == n) {
      n.parent.left = r;
    } else {
      n.parent.right = r;
    }
    r.parent = n.parent;
    n.parent = r;

    r.size = n.size;
    n.size = size(n.left) + size(n.right) + 1;
  }

  /**
   * flip the color of n, and its two direct children (must have!)
   *
   * @param n the parent of the nodes to be flipped
   */
  private void flipColors(Node n) {
    n.color = !n.color;
    n.left.color = !n.left.color;
    n.right.color = !n.right.color;
  }

  /*************************************************
   * Traversal helpers
   *************************************************/

  public void levelTraversal() {
    Node tmp;
    Queue<Node> q = new LinkedList<>();
    q.add(root);
    while (!q.isEmpty()) {
      System.out.println((tmp = q.remove()).key + (tmp.color ? "(RED)" : "(BLACK)"));
      if (tmp.left != null) q.add(tmp.left);
      if (tmp.right != null) q.add(tmp.right);
    }
  }

  public void preOrder() {
    System.out.println("===== preOrder =====");
    preOrder(root);
  }

  private void preOrder(Node n) {
    if (n == null) return;
    System.out.println(n.key + (n.color ? "(RED)" : "(BLACK)"));
    preOrder(n.left);
    preOrder(n.right);
  }
}
